package com.db_tenant;

import com.db_tenant.citus.CitusTenant;
import com.db_tenant.citus.CoordinateTenant;
import com.db_tenant.citus.WorkerTenant;
import io.vavr.CheckedFunction0;
import io.vavr.control.Try;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

@Component
public final class ReadWriteContext {

  private static class UserContextThreadLocal extends ThreadLocal<CitusTenant> {

    @Override
    protected CitusTenant initialValue() {
      return null;
    }
  }

  private static final UserContextThreadLocal LOCAL = new UserContextThreadLocal();

  //----------------
  private static ReadWriteContext RW_CONTEXT;

  public static <T> Try<T> exeInCoord(CheckedFunction0<T> function) {
    if (RW_CONTEXT != null) RW_CONTEXT.setCoord();

    try {
      return Try.of(function);
    } catch (Exception e) {
      throw e;
    } finally {
      useDefault();
    }
  }

  public static <T> Try<T> exeInWorder(CheckedFunction0<T> function) {
    if (RW_CONTEXT != null) RW_CONTEXT.setWorker();

    try {
      return Try.of(function);
    } catch (Exception e) {
      throw e;
    } finally {
      useDefault();
    }
  }

  public static void useDefault() {
    if (RW_CONTEXT != null) RW_CONTEXT.reset();
  }

  //----------------
  @Autowired
  private final CoordinateTenant coordTenant;

  @Autowired
  @Qualifier("workderTennants")
  private final List<WorkerTenant> workderTenants;

  private int workerCount = 0;
  private final AtomicInteger rr = new AtomicInteger();

  public ReadWriteContext(CoordinateTenant coordTenant, List<WorkerTenant> workderTenants) {
    this.coordTenant = coordTenant;
    this.workderTenants = workderTenants;

    workerCount = workderTenants.size();

    RW_CONTEXT = this;
  }

  //----------------
  public void setCoord() {
    LOCAL.set(coordTenant);
  }

  public void setWorker() {
    LOCAL.set(workderTenants.get(rr.incrementAndGet() % workerCount));
  }

  public CitusTenant get() {
    return LOCAL.get();
  }

  public void reset() {
    LOCAL.remove();
  }

}
